package org.movee.jalgo.search;

import lombok.extern.slf4j.Slf4j;

import java.util.ArrayList;
import java.util.AbstractMap.SimpleEntry;
import java.util.List;
import java.util.Map.Entry;

@Slf4j
public class BplusTree<K extends Comparable<K>, V> {

    // 最小出度，最大出度为2m
    protected int maxOrder;
    protected int minOrder;
    // B+树根节点
    protected Node<K, V> root;
    // 叶子节点链表头
    protected Node<K, V> head;

    public BplusTree(int maxOrder){
        if (maxOrder < 3) {
            throw new IllegalArgumentException("minimum fanout must be greater than 2");
        }
        this.maxOrder = maxOrder;
        this.minOrder = maxOrder / 2;
        root = new Node<>(true, true);
        head = root;
    }

    public Node getHead() {
        return head;
    }

    public void setHead(Node head) {
        this.head = head;
    }

    public Node getRoot() {
        return root;
    }

    public void setRoot(Node root) {
        this.root = root;
    }

    public int getMaxOrder() {
        return maxOrder;
    }

    public int getMinOrder() {
        return minOrder;
    }

    public boolean isEmpty() {
        return root.getEntries().isEmpty();
    }

    public V get(K key) {
        checkKeyValid(key);
        if (this.isEmpty() || root.mostLeftKey().compareTo(key) < 0) {
            // 空树或查找的key小于树中最小的key
            return null;
        }
        return root.get(key);
    }

    public boolean containsKey(K key) {
        checkKeyValid(key);
        if (this.isEmpty() || root.mostLeftKey().compareTo(key) < 0) {
            // 空树或查找的key小于树中最小的key
            return false;
        }
        return root.containsKey(key);
    }

    public void put(K key, V value) {
        checkKeyValid(key);
        // 因为root节点的分裂处理方法与非root节点不同，所以如果root节点已经满了，则立即分裂，同时使树的高度增加1，
        // 这样下层节点分裂需要向root节点插入key时，root节点总是有足够的空间
        maybeSplitRoot();
        // 查找key要插入的leaf node
        Node<K, V> leafNode = root.findLeafNode(key);
        // 将数据插入到leaf node中
        insertKV(key, value, leafNode);
        // 如果需要分裂，则自底向上逐层分裂
        maybeSplitNode(leafNode);
    }

    public void remove(K key) {
        // 删除的步骤
        // 1、删除目标节点中的key、value
        // 2、判断删除后，目标节点的key个数是否小于m/2
        // 3、大于或等于m/2，则判断节点的索引key是否改变，如果变了，则修改父节点中的索引值
        // 4、小于m/2，则
        // 4.1 如果左边的节点的key数量大于m/2，则将左边节点的最大key移到目标节点中，然后修改父节点中的索引值
        // 4.2 如果右边的节点的key数量大于m/2，则将右边节点的最小key移到目标节点中，然后修改右边节点父节点中的索引值
        // 4.3 左边和右边的key数量都为m/2，将目标节点的key全部合入左边节点，删除目标节点中父节点中的索引

        checkKeyValid(key);

        Node<K, V> leafNode = root.findLeafNode(key);

        leafNode.removeKey(key);

    }

    private void maybeSplitRoot() {
        if (root.size() > maxOrder) {
            Node<K, V> newNode = root.split();
            Node<K, V> newRoot = new Node<>(false, true);
            newRoot.addKey(root.mostLeftEntry());
            newRoot.addKey(newNode.mostLeftEntry());
            newRoot.addChild(root);
            newRoot.addChild(newNode);

            newNode.setParent(newRoot);
            root.setParent(newRoot);
            root.setRoot(false);

            root = newRoot;
        }
    }

    private void insertKV(K key, V value, Node<K, V> leafNode) {

        SimpleEntry<K, V> newEntry = new SimpleEntry<>(key, value);

        // 查找key在leaf node中插入的位置
        int floor = leafNode.floor(key);
        if (floor < 0) {
            // 新插入的key小于树中最小的key，修改索引值为新插入的key
            Node<K, V> parent = leafNode.getParent();
            while (parent != null && parent.mostLeftKey().compareTo(leafNode.mostLeftKey()) == 0) {
                parent.replaceKey0(newEntry);
                parent = parent.getParent();
            }
        }

        // 修改value或插入新的key，value
        if (floor >= 0 && leafNode.getKey(floor).compareTo(key) == 0) {
            leafNode.setKey(floor, newEntry);
        } else {
            leafNode.addKey(floor + 1, newEntry);
        }
    }

    private void maybeSplitNode(Node<K, V> node) {
        while (node.size() > maxOrder) {
            // 分裂节点
            Node<K, V> newNode = node.split();
            // 将新创建的节点插入到父节点中
            Node<K, V> parent = node.getParent();
            // 新节点需插入到父节点中的位置，floor肯定大于或等于0
            int floor = parent.floor(newNode.mostLeftKey());
            parent.addKey(floor + 1, newNode.mostLeftEntry());
            parent.addChild(floor + 1, newNode);

            node = parent;
        }
    }

    protected void checkKeyValid(K key) {
        if (key == null) {
            throw new IllegalArgumentException("Key must be not null!");
        }
    }

    private static class Node<K extends Comparable<K>, V> {

        // 是否为叶子节点
        protected boolean isLeaf;
        // 是否为根节点
        protected boolean isRoot;
        // 父节点
        protected Node<K, V> parent;
        // 子节点
        protected List<Node<K, V>> children;
        // 叶节点的前驱节点
        protected Node<K, V> previous;
        // 叶节点的后继节点
        protected Node<K, V> next;
        // 节点的关键字
        protected List<Entry<K, V>> entries;

        public Node(boolean isLeaf, boolean isRoot) {
            this.isLeaf = isLeaf;
            this.isRoot = isRoot;
            entries = new ArrayList<>();
            if (!isLeaf) {
                children = new ArrayList<>();
            }
        }

        public int size() {
            return entries.size();
        }

        public K getKey(int index) {
            return entries.get(index).getKey();
        }

        public K mostLeftKey() {
            return entries.get(0).getKey();
        }

        public Entry<K, V> mostLeftEntry() {
            return entries.get(0);
        }

        public void replaceKey0(Entry<K, V> entry) {
            entries.set(0, entry);
        }

        public void setKey(int index, Entry<K, V> entry) {
            entries.set(index, entry);
        }

        public void addKey(int index, Entry<K, V> entry) {
            entries.add(index, entry);
        }

        public void addKey(Entry<K, V> entry) {
            entries.add(entry);
        }

        public void addChild(int index, Node<K, V> node) {
            children.add(index, node);
        }

        public void addChild(Node<K, V> node) {
            children.add(node);
        }

        public void removeKey(K key) {
            int floor = floor(key);
            if (floor > 0 && mostLeftKey().compareTo(key) == 0) {
                entries.remove(key);
            }
        }

        /**
         * 对于非root节点，返回的floor值都会大于等于0
         * @param key key
         * @return floor值
         */
        private int floor(K key) {
            int floor = entries.size() - 1;
            for (; floor >= 0; floor--) {
                if (entries.get(floor).getKey().compareTo(key) <= 0) {
                    break;
                }
            }

            return floor;
        }

        public V get(K key) {

            int floor = floor(key);
            if (isLeaf) {
                // 已经是叶子节点，直接判断floor key是否与查找key相同
                if (entries.get(floor).getKey().compareTo(key) == 0) {
                    return entries.get(floor).getValue();
                } else {
                    return null;
                }
            } else {
                // 在子节点中递归查找
                return children.get(floor).get(key);
            }
        }

        public boolean containsKey(K key) {
            int floor = floor(key);
            if (entries.get(floor).getKey().compareTo(key) == 0) {
                return true;
            } else if (!isLeaf) {
                // 在子节点中递归查找
                return children.get(floor).containsKey(key);
            }

            return false;
        }

        /**
         * 查找key可能所在的leaf node
         * @param key key
         * @return leaf node，不会为null
         */
        public Node<K, V> findLeafNode(K key) {
            if (!isLeaf) {
                int floor = floor(key);
                floor = Math.max(0, floor);
                return children.get(floor).findLeafNode(key);
            }

            return this;
        }

        public Node<K, V> split() {

            int size = entries.size();
            int leftSize = (size + 1) / 2;

            Node<K, V> newNode = new Node<>(isLeaf, false);
            List<Entry<K, V>> rightEntries = new ArrayList<>();
            List<Node<K, V>> rightChidren = new ArrayList<>();
            for (int i = leftSize; i < size; i ++) {
                rightEntries.add(entries.get(i));
                rightChidren.add(children.get(i));
            }
            newNode.setEntries(rightEntries);
            newNode.setChildren(rightChidren);
            newNode.setParent(parent);
            newNode.setPrevious(this);
            newNode.setNext(next);
            if (next != null) {
                next.setPrevious(newNode);
            }
            next = newNode;

            List<Entry<K, V>> leftEntries = new ArrayList<>();
            List<Node<K, V>> leftChidren = new ArrayList<>();
            for (int i = 0; i < leftSize; i ++) {
                leftEntries.add(entries.get(i));
                leftChidren.add(children.get(i));
            }
            entries = leftEntries;
            children = leftChidren;

            return newNode;
        }

        public Node getPrevious() {
            return previous;
        }

        public void setPrevious(Node previous) {
            this.previous = previous;
        }

        public Node getNext() {
            return next;
        }

        public void setNext(Node next) {
            this.next = next;
        }

        public boolean isLeaf() {
            return isLeaf;
        }

        public void setLeaf(boolean isLeaf) {
            this.isLeaf = isLeaf;
        }

        public Node<K, V> getParent() {
            return parent;
        }

        public void setParent(Node<K, V> parent) {
            this.parent = parent;
        }

        public List<Entry<K, V>> getEntries() {
            return entries;
        }

        public void setEntries(List<Entry<K, V>> entries) {
            this.entries = entries;
        }

        public List<Node<K, V>> getChildren() {
            return children;
        }

        public void setChildren(List<Node<K, V>> children) {
            this.children = children;
        }

        public boolean isRoot() {
            return isRoot;
        }

        public void setRoot(boolean isRoot) {
            this.isRoot = isRoot;
        }

        public String toString(){
            StringBuilder sb = new StringBuilder();
            sb.append("isRoot: ");
            sb.append(isRoot);
            sb.append(", ");
            sb.append("isLeaf: ");
            sb.append(isLeaf);
            sb.append(", ");
            sb.append("keys: ");
            for (Entry entry : entries){
                sb.append(entry.getKey());
                sb.append(", ");
            }
            sb.append(", ");
            return sb.toString();

        }

    }

}
